from typing import Dict, Any
from ..base import BaseComponent, ComponentConfig, ComponentType
from .openai_llm import OpenAILLM

class OpenAIComponent(BaseComponent):
    """OpenAI模型组件实现"""
    
    def __init__(self, config: ComponentConfig):
        super().__init__(config)
        self.model = OpenAILLM(self.config.config)
    
    @classmethod
    def get_component_type(cls) -> ComponentType:
        return ComponentType.LLM
    
    @classmethod
    def get_component_name(cls) -> str:
        return "OpenAI模型"
    
    @classmethod
    def get_component_description(cls) -> str:
        return "使用OpenAI的API进行模型调用"
    
    @classmethod
    def get_config_schema(cls) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "system_prompt": {
                    "type": "string",
                    "title": "系统提示词",
                    "description": "模型的系统提示词"
                },
                "user_prompt": {
                    "type": "string",
                    "title": "用户提示词",
                    "description": "模型的用户提示词，可以使用{变量}的形式"
                },
                "model_name": {
                    "type": "string",
                    "title": "模型名称",
                    "description": "使用的OpenAI模型名称",
                    "default": "gpt-3.5-turbo"
                }
            },
            "required": ["model_name"]
        }
    
    @classmethod
    def get_input_schema(cls) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "input": {
                    "type": "string",
                    "title": "输入文本",
                    "description": "需要处理的文本"
                }
            },
            "required": ["input"]
        }
    
    @classmethod
    def get_output_schema(cls) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "response": {
                    "type": "string",
                    "title": "模型响应",
                    "description": "模型的输出结果"
                }
            },
            "required": ["response"]
        }
    
    async def execute(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """执行OpenAI模型调用"""
        system_prompt = self.config.config.get("system_prompt", "")
        user_prompt = self.config.config.get("user_prompt", "").format(**input_data)
        
        response = self.model.chat(
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ]
        )
        
        return {"response": response}
    
    def validate_input(self, input_data: Dict[str, Any]) -> bool:
        """验证输入数据"""
        required_fields = self.config.input_schema.get("required", [])
        return all(field in input_data for field in required_fields)
    
    def validate_output(self, output_data: Dict[str, Any]) -> bool:
        """验证输出数据"""
        return "response" in output_data